"""
散点图动画生成器
按频率f从小到大生成xy散点图，并合成GIF动画
"""

import os
import sys
import re
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import pandas as pd


def format_particle_latex(particle):
    """
    将粒子名称转换为LaTeX格式
    例如: Al+ -> Al^{+}, Al2+ -> Al^{2+}, Fe3+ -> Fe^{3+}, O2- -> O^{2-}
    """
    if not particle:
        return ''
    
    # 匹配元素符号和电荷部分 (如 Al2+, Fe3+, O2-)
    match = re.match(r'^([A-Za-z]+)(\d*[+-])$', particle)
    if match:
        element = match.group(1)
        charge = match.group(2)
        return f'$\\mathrm{{{element}}}^{{{charge}}}$ '
    else:
        # 如果格式不匹配，直接使用原文
        return f'$\\mathrm{{{particle}}}$ '

# 添加utils路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from utils.csv_hanlder import select_and_read_files


def create_scatter_animation(
    output_dir='result',
    output_gif='scatter_animation.gif',
    fps=2,
    dpi=100,
    figsize=(10, 10),
    marker_size=5,
    marker_alpha=0.6,
    particle=None,  # 粒子名称，如 'Al+', 'Fe2+' 等
    max_frames=100  # 最大帧数（自动采样）
):
    """
    读取CSV文件，按频率f生成xy散点图动画
    
    Parameters:
    -----------
    output_dir : str
        输出目录
    output_gif : str
        输出GIF文件名
    fps : int
        动画帧率
    dpi : int
        图像分辨率
    figsize : tuple
        图像尺寸
    marker_size : float
        散点大小
    marker_alpha : float
        散点透明度
    """
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 选择并读取文件
    dfs, filenames = select_and_read_files('csv')
    
    if not dfs:
        print("没有读取到有效数据")
        return
    
    # 合并所有数据
    all_data = pd.concat(dfs, ignore_index=True)
    print(f"\n合并后数据形状: {all_data.shape}")
    print(f"列名: {list(all_data.columns)}")
    
    # 检查必需的列
    required_cols = ['x', 'y', 'f']
    for col in required_cols:
        if col not in all_data.columns:
            print(f"错误: 缺少必需的列 '{col}'")
            return
    
    # 获取唯一的f值并排序
    all_unique_f = sorted(all_data['f'].unique())
    print(f"\n频率f值数量: {len(all_unique_f)}")
    print(f"频率范围: {min(all_unique_f)} ~ {max(all_unique_f)}")
    
    # 采样（如果帧数太多）
    if len(all_unique_f) > max_frames:
        # 均匀采样
        indices = np.linspace(0, len(all_unique_f) - 1, max_frames, dtype=int)
        unique_f = [all_unique_f[i] for i in indices]
        print(f"采样后帧数: {len(unique_f)} (从 {len(all_unique_f)} 帧采样)")
    else:
        unique_f = all_unique_f
    
    # 预先按f值分组数据（优化性能，避免每帧重复筛选）
    print("\n预处理数据（按f分组）...")
    grouped_data = {f: group[['x', 'y']].values for f, group in all_data.groupby('f')}
    
    # 计算全局坐标范围（用于固定坐标轴）
    x_min, x_max = all_data['x'].min(), all_data['x'].max()
    y_min, y_max = all_data['y'].min(), all_data['y'].max()
    
    # 添加边距
    x_margin = (x_max - x_min) * 0.1
    y_margin = (y_max - y_min) * 0.1
    x_range = (x_min - x_margin, x_max + x_margin)
    y_range = (y_min - y_margin, y_max + y_margin)
    
    print(f"X范围: {x_range}")
    print(f"Y范围: {y_range}")
    
    # 粒子名称（LaTeX格式）
    particle_label = format_particle_latex(particle)
    
    # 创建图形
    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    
    # 初始化散点图
    scatter = ax.scatter([], [], s=marker_size, alpha=marker_alpha, c='blue')
    
    # 设置坐标轴
    ax.set_xlim(x_range)
    ax.set_ylim(y_range)
    ax.set_xlabel('X (mm)', fontsize=12)
    ax.set_ylabel('Y (mm)', fontsize=12)
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    
    # 添加标题
    title = ax.set_title('', fontsize=14)
    
    # 添加粒子计数文本
    count_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, 
                         fontsize=10, verticalalignment='top',
                         bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
    
    def init():
        scatter.set_offsets(np.empty((0, 2)))
        title.set_text('')
        count_text.set_text('')
        return scatter, title, count_text
    
    def update(frame):
        f_value = unique_f[frame]
        
        # 使用预分组的数据（无需每帧筛选）
        xy_data = grouped_data[f_value]
        
        # 更新散点
        scatter.set_offsets(xy_data)
        
        # 更新标题和计数
        title.set_text(f'{particle_label}XY Scatter Plot (f = {f_value:.1f})')
        count_text.set_text(f'Particles: {len(xy_data)}')
        
        # 显示进度
        if (frame + 1) % 100 == 0 or frame == len(unique_f) - 1:
            print(f'  进度: {frame + 1}/{len(unique_f)} ({100*(frame+1)/len(unique_f):.1f}%)')
        
        return scatter, title, count_text
    
    # 创建动画
    duration = len(unique_f) / fps
    print(f"\n正在生成动画 ({len(unique_f)} 帧, {fps}fps, 播放时长 {duration:.2f}秒)...")
    anim = FuncAnimation(
        fig, update, init_func=init,
        frames=len(unique_f),
        interval=1000//fps,
        blit=True
    )
    
    # 保存GIF
    output_path = os.path.join(output_dir, output_gif)
    print(f"保存GIF到: {output_path}")
    
    writer = PillowWriter(fps=fps)
    anim.save(output_path, writer=writer)
    
    plt.close()
    print(f"\n动画已保存: {output_path}")
    
    return output_path


def create_scatter_with_color(
    output_dir='result',
    output_gif='scatter_color_animation.gif',
    fps=2,
    color_by='ke',  # 可以是 'ke', 'vx', 'vy', 'vz' 等
    particle=None   # 粒子名称，如 'Al+', 'Fe2+' 等
):
    """
    带颜色映射的散点图动画
    """
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    dfs, filenames = select_and_read_files('csv')
    
    if not dfs:
        print("没有读取到有效数据")
        return
    
    all_data = pd.concat(dfs, ignore_index=True)
    
    # 检查颜色列是否存在
    if color_by not in all_data.columns:
        print(f"警告: 列 '{color_by}' 不存在，使用默认颜色")
        color_by = None
    
    unique_f = sorted(all_data['f'].unique())
    
    x_min, x_max = all_data['x'].min(), all_data['x'].max()
    y_min, y_max = all_data['y'].min(), all_data['y'].max()
    x_margin = (x_max - x_min) * 0.1
    y_margin = (y_max - y_min) * 0.1
    x_range = (x_min - x_margin, x_max + x_margin)
    y_range = (y_min - y_margin, y_max + y_margin)
    
    # 颜色范围
    if color_by:
        c_min, c_max = all_data[color_by].min(), all_data[color_by].max()
    
    # 预先按f值分组数据（优化性能）
    print("\n预处理数据（按f分组）...")
    if color_by:
        grouped_data = {f: group[['x', 'y', color_by]].values for f, group in all_data.groupby('f')}
    else:
        grouped_data = {f: group[['x', 'y']].values for f, group in all_data.groupby('f')}
    
    # 粒子名称（LaTeX格式）
    particle_label = format_particle_latex(particle)
    
    fig, ax = plt.subplots(figsize=(10, 10), dpi=100)
    
    def update(frame):
        ax.clear()
        f_value = unique_f[frame]
        
        # 使用预分组的数据
        data = grouped_data[f_value]
        x_data = data[:, 0]
        y_data = data[:, 1]
        
        if color_by:
            c_data = data[:, 2]
            scatter = ax.scatter(x_data, y_data, s=5, c=c_data, 
                                cmap='viridis', vmin=c_min, vmax=c_max, alpha=0.6)
            if frame == 0:
                plt.colorbar(scatter, ax=ax, label=color_by)
        else:
            ax.scatter(x_data, y_data, s=5, alpha=0.6, c='blue')
        
        ax.set_xlim(x_range)
        ax.set_ylim(y_range)
        ax.set_xlabel('X (mm)')
        ax.set_ylabel('Y (mm)')
        ax.set_title(f'{particle_label}XY Scatter (f = {f_value:.1f}, n = {len(x_data)})')
        ax.set_aspect('equal')
        ax.grid(True, alpha=0.3)
        
        # 显示进度
        if (frame + 1) % 100 == 0 or frame == len(unique_f) - 1:
            print(f'  进度: {frame + 1}/{len(unique_f)} ({100*(frame+1)/len(unique_f):.1f}%)')
    
    print(f"\n正在生成彩色动画 ({len(unique_f)} 帧)...")
    anim = FuncAnimation(fig, update, frames=len(unique_f), interval=1000//fps)
    
    output_path = os.path.join(output_dir, output_gif)
    anim.save(output_path, writer=PillowWriter(fps=fps))
    plt.close()
    
    print(f"\n动画已保存: {output_path}")
    return output_path


if __name__ == '__main__':
    import argparse
    
    # GIF格式最大fps约50
    FPS = 50
    
    parser = argparse.ArgumentParser(description='生成XY散点图动画')
    parser.add_argument('--output', type=str, default='scatter_animation.gif', help='输出文件名')
    parser.add_argument('--color', type=str, default=None, help='颜色映射列 (如 ke, vx)')
    parser.add_argument('--size', type=float, default=5, help='散点大小')
    parser.add_argument('--particle', type=str, default=None, help='粒子名称 (如 Al+, Fe2+)')
    parser.add_argument('--max-frames', type=int, default=200, help='最大帧数 (默认200，自动采样)')
    
    args = parser.parse_args()
    
    if args.color:
        create_scatter_with_color(
            output_gif=args.output,
            fps=FPS,
            color_by=args.color,
            particle=args.particle
        )
    else:
        create_scatter_animation(
            output_gif=args.output,
            fps=FPS,
            marker_size=args.size,
            particle=args.particle,
            max_frames=args.max_frames
        )

